#!/usr/bin/env python3
"""
Elasticsearch推荐查询DSL设计
支持标签匹配、文本相似、自定义打分排序
"""

from typing import List, Dict, Any, Optional
from datetime import datetime, timedelta
import json

class RecommendationDSL:
    """推荐查询DSL生成器"""
    
    def __init__(self):
        self.index_name = "content_recommend"
    
    def build_tag_matching_query(
        self, 
        user_tag_ids: List[int], 
        tag_weights: Optional[Dict[int, float]] = None,
        size: int = 20,
        from_: int = 0
    ) -> Dict[str, Any]:
        """
        构建标签匹配推荐查询
        
        Args:
            user_tag_ids: 用户兴趣标签ID列表
            tag_weights: 标签权重字典 {tag_id: weight}
            size: 返回数量
            from_: 起始位置
        """
        
        # 构建标签权重查询
        should_clauses = []
        for tag_id in user_tag_ids:
            weight = tag_weights.get(tag_id, 1.0) if tag_weights else 1.0
            
            # tag_ids是数组字段，使用term查询每个元素
            should_clauses.append({
                "term": {
                    "tag_ids": {
                        "value": tag_id,
                        "boost": weight
                    }
                }
            })
        
        query = {
            "index": self.index_name,
            "body": {
                "query": {
                    "bool": {
                        "must": [
                            {
                                "term": {
                                    "status": 2  # 已发布
                                }
                            }
                        ],
                        "should": should_clauses,
                        "minimum_should_match": 1
                    }
                },
                "size": size,
                "from": from_,
                "sort": [
                    "_score",
                    {"published_at": {"order": "desc"}}
                ]
            }
        }
        
        return query
    
    def build_text_similarity_query(
        self,
        content_id: int,
        fields: List[str] = None,
        size: int = 20,
        min_score: float = 0.1
    ) -> Dict[str, Any]:
        """
        构建文本相似推荐查询
        
        Args:
            content_id: 参考内容ID
            fields: 相似度计算字段
            size: 返回数量
            min_score: 最小相似度分数
        """
        
        if fields is None:
            fields = ["title", "description"]
        
        query = {
            "index": self.index_name,
            "body": {
                "query": {
                    "bool": {
                        "must": [
                            {
                                "term": {
                                    "status": 2  # 已发布
                                }
                            },
                            {
                                "more_like_this": {
                                    "fields": fields,
                                    "like": [
                                        {
                                            "_index": self.index_name,
                                            "_id": content_id
                                        }
                                    ],
                                    "min_term_freq": 1,
                                    "min_doc_freq": 1,
                                    "max_query_terms": 12,
                                    "minimum_should_match": "30%"
                                }
                            }
                        ],
                        "must_not": [
                            {
                                "term": {
                                    "content_id": content_id  # 排除自身
                                }
                            }
                        ]
                    }
                },
                "size": size,
                "min_score": min_score,
                "sort": [
                    "_score",
                    {"published_at": {"order": "desc"}}
                ]
            }
        }
        
        return query
    
    def build_custom_scoring_query(
        self,
        base_query: Dict[str, Any],
        boost_factors: Optional[Dict[str, float]] = None
    ) -> Dict[str, Any]:
        """
        构建自定义打分排序查询
        
        Args:
            base_query: 基础查询
            boost_factors: 各因子权重
        """
        
        if boost_factors is None:
            boost_factors = {
                "hot_score": 0.3,      # 热度权重
                "time_score": 0.2,     # 时间权重
                "interaction": 0.3,    # 互动权重
                "original_score": 0.2  # 原始分数权重
            }
        
        # 构建function_score查询
        function_score_query = {
            "index": self.index_name,
            "body": {
                "query": {
                    "function_score": {
                        "query": base_query["body"]["query"],
                        "functions": [
                            # 热度分数
                            {
                                "field_value_factor": {
                                    "field": "hot_score",
                                    "factor": boost_factors["hot_score"],
                                    "modifier": "log1p"
                                }
                            },
                            # 时间分数
                            {
                                "field_value_factor": {
                                    "field": "time_score", 
                                    "factor": boost_factors["time_score"],
                                    "modifier": "log1p"
                                }
                            },
                            # 互动数据综合分数
                            {
                                "script_score": {
                                    "script": {
                                        "source": """
                                        double interaction_score = 0;
                                        interaction_score += Math.log1p(doc['like_count'].value) * 0.4;
                                        interaction_score += Math.log1p(doc['comment_count'].value) * 0.3;
                                        interaction_score += Math.log1p(doc['share_count'].value) * 0.2;
                                        interaction_score += Math.log1p(doc['view_count'].value) * 0.1;
                                        return interaction_score * params.factor;
                                        """,
                                        "params": {
                                            "factor": boost_factors["interaction"]
                                        }
                                    }
                                }
                            }
                        ],
                        "score_mode": "sum",
                        "boost_mode": "multiply",
                        "min_score": 0.1
                    }
                },
                "size": base_query["body"].get("size", 20),
                "from": base_query["body"].get("from", 0),
                "sort": [
                    {"_score": {"order": "desc"}},
                    {"published_at": {"order": "desc"}}
                ]
            }
        }
        
        return function_score_query
    
    def build_hybrid_recommendation_query(
        self,
        user_tag_ids: List[int],
        reference_content_id: Optional[int] = None,
        tag_weights: Optional[Dict[int, float]] = None,
        size: int = 20,
        from_: int = 0
    ) -> Dict[str, Any]:
        """
        构建混合推荐查询（标签匹配 + 文本相似）
        
        Args:
            user_tag_ids: 用户兴趣标签
            reference_content_id: 参考内容ID（用于文本相似）
            tag_weights: 标签权重
            size: 返回数量
            from_: 起始位置
        """
        
        # 构建标签匹配查询
        tag_query = self.build_tag_matching_query(
            user_tag_ids, tag_weights, size=size, from_=from_
        )
        
        # 如果有参考内容，添加文本相似查询
        if reference_content_id:
            text_query = self.build_text_similarity_query(
                reference_content_id, size=size
            )
            
            # 混合查询：标签匹配 OR 文本相似
            hybrid_query = {
                "index": self.index_name,
                "body": {
                    "query": {
                        "bool": {
                            "must": [
                                {
                                    "term": {
                                        "status": 2  # 已发布
                                    }
                                }
                            ],
                            "should": [
                                tag_query["body"]["query"],  # 标签匹配
                                text_query["body"]["query"]  # 文本相似
                            ],
                            "minimum_should_match": 1
                        }
                    },
                    "size": size,
                    "from": from_,
                    "sort": [
                        "_score",
                        {"published_at": {"order": "desc"}}
                    ]
                }
            }
            
            return hybrid_query
        
        return tag_query
    
    def build_search_after_query(
        self,
        user_tag_ids: List[int],
        search_after: Optional[List] = None,
        tag_weights: Optional[Dict[int, float]] = None,
        size: int = 20
    ) -> Dict[str, Any]:
        """
        构建search_after分页查询（避免深分页问题）
        
        Args:
            user_tag_ids: 用户兴趣标签
            search_after: 上一页的排序值
            tag_weights: 标签权重
            size: 返回数量
        """
        
        query = self.build_tag_matching_query(
            user_tag_ids, tag_weights, size=size, from_=0
        )
        
        # 移除from参数，使用search_after
        query["body"].pop("from", None)
        
        if search_after:
            query["body"]["search_after"] = search_after
        
        return query

# 查询DSL示例
class QueryExamples:
    """查询DSL示例"""
    
    @staticmethod
    def get_tag_matching_example():
        """标签匹配查询示例"""
        dsl_builder = RecommendationDSL()
        
        # 用户兴趣标签：[1,3,5]，权重：美食(2.0)、旅行(1.5)、咖啡(1.0)
        user_tags = [1, 3, 5]
        tag_weights = {1: 2.0, 3: 1.5, 5: 1.0}
        
        query = dsl_builder.build_tag_matching_query(
            user_tag_ids=user_tags,
            tag_weights=tag_weights,
            size=20
        )
        
        return query
    
    @staticmethod
    def get_text_similarity_example():
        """文本相似查询示例"""
        dsl_builder = RecommendationDSL()
        
        # 基于content_id=123的内容推荐相似内容
        query = dsl_builder.build_text_similarity_query(
            content_id=123,
            fields=["title", "description"],
            size=20,
            min_score=0.1
        )
        
        return query
    
    @staticmethod
    def get_custom_scoring_example():
        """自定义打分排序示例"""
        dsl_builder = RecommendationDSL()
        
        # 基础标签匹配查询
        base_query = dsl_builder.build_tag_matching_query([1, 3, 5])
        
        # 自定义打分权重
        boost_factors = {
            "hot_score": 0.3,      # 热度权重30%
            "time_score": 0.2,     # 时间权重20%
            "interaction": 0.3,    # 互动权重30%
            "original_score": 0.2  # 原始分数权重20%
        }
        
        query = dsl_builder.build_custom_scoring_query(
            base_query, boost_factors
        )
        
        return query

if __name__ == "__main__":
    # 测试DSL生成
    examples = QueryExamples()
    
    print("=== 标签匹配查询示例 ===")
    print(json.dumps(examples.get_tag_matching_example(), indent=2, ensure_ascii=False))
    
    print("\n=== 文本相似查询示例 ===")
    print(json.dumps(examples.get_text_similarity_example(), indent=2, ensure_ascii=False))
    
    print("\n=== 自定义打分排序示例 ===")
    print(json.dumps(examples.get_custom_scoring_example(), indent=2, ensure_ascii=False))
